# 预处理
import os
import mne
from scipy.io import loadmat
import numpy as np


def bandpass_filter(data, sfreq, low_freq, high_freq):
    filtered_data = mne.filter.filter_data(
        data,
        sfreq=sfreq,
        l_freq=low_freq,
        h_freq=high_freq,
        method='fir',
        verbose=False
    )
    return filtered_data


def seg_raw_to_all_bands(all_bands, root_dir, sfreq):
    # 把sxx.dat等数据分割成不同的频段的数据
    baseline_seconds = 3
    baseline_samples = baseline_seconds * sfreq

    for sub_file in os.listdir(f'{root_dir}/data_preprocessed_matlab'):
        if (sub_file.startswith('s')):
            print(sub_file)
            for band_name, band_range in all_bands.items():
                sub = loadmat(
                    f'{root_dir}/data_preprocessed_matlab/{sub_file}')
                data_without_baseline = sub['data'][:, :32, baseline_samples:]
                data_filtered = bandpass_filter(
                    data_without_baseline, sfreq, band_range[0], band_range[1])

                file_path = f'{root_dir}/all_bands_data/{band_name}'
                os.makedirs(file_path, exist_ok=True)

                np.save(
                    f'{file_path}/{os.path.splitext(sub_file)[0]}.npy', data_filtered)

                labels = sub['labels']
                np.save(
                    f'{file_path}/{os.path.splitext(sub_file)[0]}_labels.npy', labels)


if __name__ == "__main__":
    all_bands = {
        "delta": (0.1, 3),
        "theta": (4, 7),
        "low-alpha": (8, 9.5),
        "high-alpha": (10.5, 12),
        "alpha": (8, 12),
        "low_beta": (13, 16),
        "mid-beta": (17, 20),
        "high-beta": (21, 29),
        "beta": (13, 29),
        "gamma": (30, 50)
    }
    sfreq = 128
    root_dir = 'C:/Users/john/Desktop/learn-eeg'

    seg_raw_to_all_bands(all_bands, root_dir, sfreq)
